package com.qdhh.enrollment.config.datasource;

import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

/**
 * Routes read-only traffic to the replica datasource while keeping writes on the master.
 */
@Aspect
@Component
@Order(Ordered.HIGHEST_PRECEDENCE + 1)
public class ReadWriteDataSourceAspect {

    private static final String SLAVE_DATASOURCE_KEY = "slave";

    @Around("@annotation(readOnly)")
    public Object routeWithReadOnlyAnnotation(ProceedingJoinPoint joinPoint, ReadOnly readOnly) throws Throwable {
        return executeAgainstReplica(joinPoint);
    }

    @Around("@annotation(transactional)")
    public Object routeTransactionalReadOnly(ProceedingJoinPoint joinPoint, Transactional transactional) throws Throwable {
        if (transactional.readOnly()) {
            return executeAgainstReplica(joinPoint);
        }
        return joinPoint.proceed();
    }

    private Object executeAgainstReplica(ProceedingJoinPoint joinPoint) throws Throwable {
        DynamicDataSourceContextHolder.push(SLAVE_DATASOURCE_KEY);
        try {
            return joinPoint.proceed();
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }
}
